package com.feheadline.ml.classification;

import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectOutputStream;
import java.util.List;

import org.apache.commons.lang3.exception.ExceptionUtils;
import org.apache.log4j.Logger;

import com.aliasi.classify.Classification;
import com.aliasi.classify.Classified;
import com.aliasi.classify.TfIdfClassifierTrainer;
import com.aliasi.tokenizer.TokenFeatureExtractor;
import com.dtruth.dataflow.dao.ClassifierDAO;

/**
 * 
 * @author 岳海亮
 * @email yhl@feheadline.com
 * @date 2014年10月28日
 */
public class TrainClassifier {
	private static final Logger logger = Logger
			.getLogger(TrainClassifier.class);
	private String modelFile;
	private String tableName;

	public TrainClassifier(String tableName, String modelFile) {
		this.modelFile = modelFile;
		this.tableName = tableName;
	}

	public void train() {
		long start = System.currentTimeMillis();
		TfIdfClassifierTrainer<CharSequence> classifier = new TfIdfClassifierTrainer<CharSequence>(
				new TokenFeatureExtractor(AnsjTokenizerFactory.INSTANCE));
		List<String> channels = getAllChannels(tableName);
		for (String channel : channels) {
			Classification classification = new Classification(channel);
			String[] _c = channel.split("@");
			List<String> data = null;
			if (_c.length > 1)
				data = getDataByChannel(tableName, _c[0], _c[1]);
			else
				data = getDataByChannel(tableName, _c[0], "");
			System.out.println(channel + ":" + data.size());
			for (String text : data) {
				Classified<CharSequence> classified = new Classified<CharSequence>(
						text, classification);
				classifier.handle(classified);
			}
		}
		System.out.println("Save model to " + modelFile);
		ObjectOutputStream os;
		try {
			os = new ObjectOutputStream(new FileOutputStream(modelFile));
			classifier.compileTo(os);
			os.close();
		} catch (IOException e) {
			logger.error(ExceptionUtils.getStackTrace(e));
		}

		long end = System.currentTimeMillis();
		System.out.println("time cost:" + (end - start) + "ms");
	}

	public List<String> getDataByChannel(String tableName, String channel,
			String subchannel) {
		ClassifierDAO cd = new ClassifierDAO();
		List<String> list = cd.listByChannel(tableName, channel, subchannel);
		cd.close();
		return list;
	}

	public List<String> getAllChannels(String tableName) {
		ClassifierDAO cd = new ClassifierDAO();
		List<String> list = cd.getAllChannels(tableName);
		cd.close();
		return list;
	}
}
